# Responsability
# Gelman Rubin statistics and storage of the GR markov chain data

GelmanRubin <- setRefClass(    
  "gelmanrubin"
  
  , fields = list(
      chainsParameterSums="list",
      iteration="numeric",
      convergenceConst="numeric",
      burnPhaseIterations="numeric",
      RHAT="numeric", # Rhat Is Used As Break condition, after convergence of chains, rhat is reduced by 1 in each iteration until it becomes 0
      SAMPLE_NO="numeric"
  )
  , methods = list(
    #
    #
    # Constructor
    #
    #
    initialize = function(..., chainsParameterSums=list(), iteration=0, 
                          convergenceConst=1.1, burnPhaseIterations=2, 
                          SAMPLE_NO=5,RHAT=0, ids=list())
    {
      loginfo("Initialize GelmanRubin")
      
      .self$RHAT <<- SAMPLE_NO +1
      
      .self$chainsParameterSums <<- list()
      for (i in 1:length(ids)) {
        loginfo(ids[i])
        .self$chainsParameterSums[[as.numeric(ids[i])]] <<- list(ji=NULL, jiSquared=NULL)
        
      }
        
      
      #chainsParameterSums <- list(sumji=list(), sumjiSquared=list())
      #For access chainsParametersSums$sumji[[i]] <- el dataframe de parameters
      
      callSuper(..., 
                iteration=iteration,
                convergenceConst=convergenceConst,
                burnPhaseIterations=burnPhaseIterations,
                SAMPLE_NO=SAMPLE_NO)
      
    },
    #
    #
    # New Iteration. Update chain parameters sum and iteration counter
    #
    #
    update = function(chainState) {

        logdebug("METHOD IN: GelmanRubin$update")
        print(chainState)
        
        if (is.null(chainsParameterSums[[chainState$id]]$ji)) {
          # sum up Parameters and squared Parameters s_j = 1/(n-1) * (  sumji_squared  - 1/n * ( sumji )^2 )
          chainsParameterSums[[chainState$id]]$ji <<- chainState$lastParameters
          chainsParameterSums[[chainState$id]]$jiSquared <<- (chainState$lastParameters)^2
        } else {
          chainsParameterSums[[chainState$id]]$ji <<- chainsParameterSums[[chainState$id]]$ji + chainState$lastParameters
          chainsParameterSums[[chainState$id]]$jiSquared <<- chainsParameterSums[[chainState$id]]$jiSquared + (chainState$lastParameters)^2
        }
      
        iteration <<- chainState$iteration
        
        logdebug("METHOD OUT: GelmanRubin$update")
    },
    #
    #
    # Gelman Rubin convergence diagnostic. Return TRUE when the chains converge after the burn phase
    #
    #
    convergenceDiagnostic = function() {

      logdebug("METHOD IN: GelmanRubin$convergenceDiagnostic")
      
      convergence = FALSE
      
      if ((length(chainsParameterSums) > 1) & ( iteration > 2 ) )
      {
        # calculate W and B see Gelman Rubin Convergence
       
        withinVariance <- as.data.frame(lapply( X=chainsParameterSums, .self$calcWithinVariance))  # matrix of variances of all Parameters k x (length(chain) = m )
        # calculate mean of all parameters over all chains
        W <- apply( withinVariance , 1, mean )
    
        withinMean <- as.data.frame(lapply( X=chainsParameterSums, .self$calcWithinMean ))  # matrix of means of all Parameter k x (length = m )
        B <- iteration *  apply ( withinMean , 1, var )
        
        # B/i := apply() * i / i  could be canceled of course, inserted just to follow theory
        varianceHat <- W*(iteration-1)/iteration  + B/iteration
        
        # final measure R_Hat
        measure <- sqrt(varianceHat/W)
        
        #TODO: take out the burnPhaseIterations check for skip this diagnostic during the burn phase
        convergence <- all(measure <= convergenceConst) & iteration > burnPhaseIterations
      }
        
      logdebug("METHOD OUT: GelmanRubin$convergenceDiagnostic")
      
      convergence
    },
    reachTheEnd = function() {
      
      loginfo("METHOD IN: GelmanRubin$reachTheEnd")
      
      convergence <- convergenceDiagnostic()
      loginfo(convergence)
      loginfo(RHAT)
      loginfo(SAMPLE_NO)
      # sample SAMPLE_NO points after burn in phase and at least 1000 iterations
      if (RHAT > SAMPLE_NO &  convergence)
      {
        RHAT <<- RHAT - 1
      }
      # reduce rhat by 1 after burn in phase is reached
      if (RHAT <= SAMPLE_NO)
      {
        RHAT <<- RHAT - 1
      }
      
      loginfo("METHOD OUT: GelmanRubin$reachTheEnd")
      
      (RHAT <= 0)
    },
    #
    #
    # 
    #
    #
    calcWithinVariance = function(chainParameterSums) {
      sj_squared <- (1/(iteration-1)) * ( chainParameterSums$jiSquared - (1/iteration) * ( chainParameterSums$ji )^2 )
      sj_squared
      
    },
    #
    #
    #
    #
    #
    calcWithinMean = function(chainParameterSums)
    {
      meanj <- 1/iteration * chainParameterSums$ji
      meanj
    }
    
  )
)